import logging
import time
from itertools import chain

import torch
import torch.nn as nn
import torch.optim as optim
from  tqdm import tqdm
# classification_report可以导出字典格式，修改参数：output_dict=True，可以将字典在保存为csv格式输出
from sklearn.metrics import precision_score, recall_score, f1_score, classification_report

from P03_NER.LSTM_CRF.config import Config
from P03_NER.LSTM_CRF.model.BiLSTM import NERLSTM
from P03_NER.LSTM_CRF.model.BiLSTM_CRF import NERLSTM_CRF
from P03_NER.LSTM_CRF.utils.data_loader import get_data, word2id

conf = Config()

# 模型训练函数
def model2dev(dev_dataloader, model, criterion=None):
    '''
    使用验证集，评估模型的效果【同时支持 NERLSTM 和 NERLSTM_CRF】
    :param dev_dataloader: 验证集dataloader
    :param model: 需要评估的模型实例
    :param criterion: 损失函数对象，因为 NERLSTM_CRF不需要损失函数对象，所以设置默认值为None
    :return:
    '''
    # 1.定义打印日志参数
    # 保存平均损失
    avg_loss = 0
    # 保存预测结果和真实标签
    preds, golds = [], []
    # 2.将模型设置为评估模式
    model.eval()
    # 3.内部遍历数据迭代器dataloader
    for index, (input_ids_padded, labels_padded, attention_mask) in enumerate(tqdm(dev_dataloader)):
        # 3.1 将数据送入模型得到输出结果
        # 把数据放到gpu上
        input_ids_padded = input_ids_padded.to(conf.device)
        labels_padded = labels_padded.to(conf.device)
        attention_mask = attention_mask.to(conf.device)
        # 送入模型
        if model.name == "BiLSTM":
            output = model(input_ids_padded, attention_mask)
            # print(f'output--->{output.shape}')  # [8, 55, 12]
            # 3.2 计算损失
            # 形状转变
            output_for_loss = output.view(-1, len(conf.tag2id))  # [8*55, 12]
            labels_for_loss = labels_padded.view(-1)  # [8*55]
            # 在计算交叉熵损失时，需要第一个参数的shape为二维的，形状是[批次大小,标签的数量]；第二个参数的shape为一维的，即标签的数量
            loss = criterion(output_for_loss, labels_for_loss)
            # print(f'loss--->{loss.item()}')
            # 对平均损失进行累加
            avg_loss += loss.item()
            # 3.3 处理结果
            predict = output.argmax(dim=-1).tolist()
            # print(f'predict--->{predict}')  # [8, 55]
        elif model.name == "BiLSTM_CRF":
            # 3.2 计算损失
            loss = model.log_likelihood(input_ids_padded, labels_padded, attention_mask).mean()
            # 对平均损失进行累加
            avg_loss += loss.item()

            # 3.3 处理结果
            predict = model(input_ids_padded, attention_mask)
            # print(f'predict--->{predict}')

        # 获取真实样本的句子长度（忽略pad的部分）
        # # 方法一：通过input_ids_padded来获取真实的句子长度，思路就是将真实的id放到一个列表中，然后获取列表的长度，即为真实的句子长度
        # real_x = [] # 用来存储真实的id
        # for sample in input_ids_padded: # 遍历每个样本
        #     temp = []  # 存储单个样本
        #     for id in sample:
        #         if id != 0:
        #             temp.append(id.item())
        #     real_x.append(temp)  # 将一个样本id放到批次的列表中
        # print(f'real_x--->{real_x}')
        #
        # # 根据真实的句子长度，获取预测的标签
        # for index, lable in enumerate(predict):
        #     # 真实的句子长度
        #     length = len(real_x[index])
        #     preds.extend(lable[:length])
        # print(f'preds--->{preds}')
        # # 根据真实的句子长度，获取真实的标签
        # for index, lable in enumerate(labels_padded.tolist()):
        #     # 真实的句子长度
        #     length = len(real_x[index])
        #     golds.extend(lable[:length])

        # # 方法二：通过input_ids_padded进行非零判断，然后进行求和，来获取每个句子的真实长度
        # real_len = (input_ids_padded>0).sum(dim=1).tolist()
        # # print(f'real_len--->{real_len}')
        # # 根据真实的句子长度，获取预测的标签
        # for index, lable in enumerate(predict):
        #     preds.extend(lable[:real_len[index]])
        # # print(f'preds--->{preds}')
        # # 根据真实的句子长度，获取真实的标签
        # for index, lable in enumerate(labels_padded.tolist()):
        #     golds.extend(lable[:real_len[index]])

        # 方法三：使用bool类型的tensor，直接获取想要的数据
        if model.name == "BiLSTM":
            predict_result = torch.tensor(predict, device=conf.device)[(input_ids_padded > 0)].tolist()
            # print(f'predict_result--->{predict_result}')
            preds.extend(predict_result)
            real_result = labels_padded[(input_ids_padded > 0)].tolist()
            # print(f'real_result--->{real_result}')
            golds.extend(real_result)
        elif model.name == "BiLSTM_CRF":
            preds.extend(chain(*predict))
            real_result = labels_padded[(input_ids_padded > 0)].tolist()
            # print(f'real_result--->{real_result}')
            golds.extend(real_result)

        # 3.4 统计批次内指标(略)
        # break

    # 4.统计整体指标
    avg_loss = avg_loss / len(dev_dataloader)
    # print(f'golds--->{golds}')
    # print(f'preds--->{preds}')
    precision = precision_score(golds, preds, average='weighted')
    recall = recall_score(golds, preds, average='weighted')
    f1 = f1_score(golds, preds, average='weighted')
    report = classification_report(golds, preds)

    return  precision, recall, f1, report, avg_loss


def model2train():
    # 1.构建数据迭代器Dataloader(包括数据处理与构建数据源Dataset)
    train_dataloader, dev_dataloader = get_data()
    # 2.实例化模型（需求进行选择）
    models = {"BiLSTM": NERLSTM,
              "BiLSTM_CRF": NERLSTM_CRF}
    model = models[conf.model](conf.embedding_dim, conf.hidden_dim, conf.dropout, word2id, conf.tag2id).to(conf.device)
    print(f'model--》{model}')
    # 3.实例化损失函数对象
    # 忽略标签为11的索引，即[PAD]，这个标签不在参与损失计算中（不对其计算损失，也不产生梯度）【原因：padding部分并不是真实的标签，不应该影响训练】
    criterion = nn.CrossEntropyLoss(ignore_index=11)
    # 4.实例化优化器对象
    optimizer = optim.Adam(model.parameters(), lr=conf.lr)
    # 5.定义打印日志参数
    start_time = time.time()

    # 6.开始训练
    f1_score = -100
    if conf.model == "BiLSTM":
        # 6.1 实现外层大循环epoch
        for epoch in range(conf.epochs):
            # 6.2 将模型设置为训练模式
            model.train()
            # 6.3 内部遍历数据迭代器dataloader
            for index, (input_ids_padded, labels_padded, attention_mask) in enumerate(tqdm(train_dataloader)):
                # 1）将数据送入模型得到输出结果
                # print(f'input_ids_padded--->{input_ids_padded.shape}')  # [8, 59]
                # print(f'labels_padded--->{labels_padded.shape}')  # [8, 59]
                # print(f'attention_mask--->{attention_mask.shape}')  # [8, 59]
                # 把数据放到gpu上
                input_ids_padded = input_ids_padded.to(conf.device)
                labels_padded = labels_padded.to(conf.device)
                attention_mask = attention_mask.to(conf.device)
                # 送入模型
                output = model(input_ids_padded, attention_mask)
                # print(f'output--->{output.shape}')  # [8, 59, 12]
                # 2）计算损失
                # 形状转变
                output = output.view(-1, len(conf.tag2id))  # [8*59, 12]
                # print(f'output--->{output.shape}')
                labels_padded = labels_padded.view(-1)  # [8*59]
                # print(f'labels_padded--->{labels_padded.shape}')
                # 在计算交叉熵损失时，需要第一个参数的shape为二维的，形状是[批次大小,标签的数量]；第二个参数的shape为一维的，即标签的数量
                loss = criterion(output, labels_padded)

                # 3）梯度清零: optimizer.zero_grad()
                optimizer.zero_grad()
                # 4）反向传播(计算梯度): loss.backward()
                loss.backward()
                # 5）梯度更新(参数更新): optimizer.step()
                optimizer.step()
                # 6）打印内部训练日志
                if (index+1) % 200 == 0:
                    print(f'epoch:{epoch+1},loss:{loss.item():.4f},time:{time.time()-start_time:.2f}')
                    # break
            # 6.4 使用验证集进行模型评估【将模型设置为评估模式】
            precision, recall, f1, report, avg_loss = model2dev(dev_dataloader, model, criterion)
            # 6.5 保存模型: torch.save(model.state_dict(), "model_path")
            if f1 > f1_score:
                f1_score = f1
                # 保存模型
                torch.save(model.state_dict(), "save_model/bilstm_best.pth")
                print(report)
    elif conf.model == "BiLSTM_CRF":
        # 6.1 实现外层大循环epoch
        for epoch in range(conf.epochs):
            # 6.2 将模型设置为训练模式
            model.train()
            # 6.3 内部遍历数据迭代器dataloader
            for index, (input_ids_padded, labels_padded, attention_mask) in enumerate(tqdm(train_dataloader)):
                # 1）将数据送入模型得到输出结果
                # print(f'input_ids_padded--->{input_ids_padded.shape}')  # [8, 59]
                # print(f'labels_padded--->{labels_padded.shape}')  # [8, 59]
                # print(f'attention_mask--->{attention_mask.shape}')  # [8, 59]
                # 把数据放到gpu上
                input_ids_padded = input_ids_padded.to(conf.device)
                labels_padded = labels_padded.to(conf.device)
                attention_mask = attention_mask.to(conf.device)
                # 2）计算损失
                # 直接调用模型的log_likelihood方法，获取模型的损失【默认是返回每个样本的损失，需要将这些损失转成标量，一般使用mean()求平均】
                loss = model.log_likelihood(input_ids_padded, labels_padded, attention_mask).mean()
                # print(f'loss--->{loss}')
                # 3）梯度清零: optimizer.zero_grad()
                optimizer.zero_grad()
                # 4）反向传播(计算梯度): loss.backward()
                loss.backward()
                # 5）梯度更新(参数更新): optimizer.step()
                optimizer.step()
                # 6）打印内部训练日志
                if (index + 1) % 200 == 0:
                    print(f'epoch:{epoch + 1},loss:{loss.item():.4f},time:{time.time() - start_time:.2f}')
                # break
            # 6.4 使用验证集进行模型评估【将模型设置为评估模式】
            precision, recall, f1, report, avg_loss = model2dev(dev_dataloader, model, criterion)
            # 6.5 保存模型: torch.save(model.state_dict(), "model_path")
            if f1 > f1_score:
                f1_score = f1
                # 保存模型
                torch.save(model.state_dict(), "save_model/bilstm_crf_best.pth")
                print(report)
    # 6.6 打印外部训练日志
    end_time = time.time()
    print(f'训练时间:{end_time-start_time:.2f}')


if __name__ == '__main__':
    # 真正训练前，需要将DataLoader中的shuffle设置为True
    model2train()